#include "types.h"
#include "param.h"
#include "spinlock.h"
#include "riscv.h"
#include "defs.h"
#include "string.h"
#include "kalloc.h"
#include "proc.h"
#include "printf.h"

// 进程表与 CPU 表
struct proc proc[NPROC];
struct cpu cpus[NCPU];

// PID 分配
static struct spinlock pid_lock;
static int nextpid = 1;

// 父子关系/等待
struct spinlock wait_lock;

extern void swtch(struct context*, struct context*);

// 本机只有 1 个 HART，cpuid 恒为 0；保留接口以便扩展 SMP
int
cpuid(void)
{
  return 0;
}

struct cpu*
mycpu(void)
{
  return &cpus[cpuid()];
}

static inline int
holding_spin(struct spinlock *lk)
{
  // 简化：沿用 spinlock.c 的 holding 实现
  extern int holding(struct spinlock *);
  return holding(lk);
}

struct proc*
myproc(void)
{
  push_off();
  struct cpu *c = mycpu();
  struct proc *p = c->proc;
  pop_off();
  return p;
}

static int
allocpid(void)
{
  int pid;
  acquire(&pid_lock);
  pid = nextpid++;
  release(&pid_lock);
  return pid;
}

static void
proc_reset(struct proc *p)
{
  p->state = UNUSED;
  p->chan = 0;
  p->killed = 0;
  p->xstate = 0;
  p->pid = 0;
  p->parent = 0;
  p->trapframe = 0;
  p->entry = 0;
  p->name[0] = 0;
  // 不释放 context 与 kstack 在此处
}

void
procinit(void)
{
  initlock(&pid_lock, "pid");
  initlock(&wait_lock, "wait_lock");
  for (int i = 0; i < NPROC; i++) {
    initlock(&proc[i].lock, "proc");
    acquire(&proc[i].lock);
    proc_reset(&proc[i]);
    release(&proc[i].lock);
  }
  for (int i = 0; i < NCPU; i++) {
    cpus[i].proc = 0;
    memset(&cpus[i].context, 0, sizeof(cpus[i].context));
    cpus[i].noff = 0;     // 新增初始化
    cpus[i].intena = 0;   // 新增初始化
  }
}

// 分配新进程结构（内核线程）
struct proc*
alloc_process(void)
{
  struct proc *p;

  for (p = proc; p < &proc[NPROC]; p++) {
    acquire(&p->lock);
    if (p->state == UNUSED) {
      p->state = USED;
      p->pid = allocpid();
      // 为内核栈分配一页
      void *kstack = kalloc();
      if (kstack == 0) {
        p->state = UNUSED;
        p->pid = 0;
        release(&p->lock);
        return 0;
      }
      p->kstack = (uint64)kstack;

      // 为trapframe分配内存（用于系统调用）
      p->trapframe = (struct trapframe *)kalloc();
      if (p->trapframe == 0) {
        kfree(kstack);
        p->state = UNUSED;
        p->pid = 0;
        release(&p->lock);
        return 0;
      }
      memset(p->trapframe, 0, sizeof(struct trapframe));

      // 初始化 context：首次 swtch 切入时将跳到 forkret_kthread
      memset(&p->context, 0, sizeof(p->context));
      p->context.sp = p->kstack + PGSIZE;
      extern void forkret_kthread(void);
      p->context.ra = (uint64)forkret_kthread;

      release(&p->lock);
      return p;
    }
    release(&p->lock);
  }
  return 0;
}

// 在 free_process() 中释放 trapframe
void
free_process(struct proc *p)
{
  if (!holding_spin(&p->lock)) {
    panic("free_process without p->lock");
  }
  if (p->kstack) {
    kfree((void*)p->kstack);
    p->kstack = 0;
  }
  if (p->trapframe) {
    kfree((void*)p->trapframe);
    p->trapframe = 0;
  }
  proc_reset(p);
}

// 内核线程第一次被调度时运行此函数
static void
kthread_entry(void)
{
  struct proc *p = myproc();

  // 检查是否被杀死
  if (p && p->killed) {
    exit_process(-1);
  }

  if (p->entry)
    p->entry();

  // 用户线程返回即退出
  exit_process(0);
}

// 与 xv6 一致：第一次切入时运行 forkret，再转入具体入口
void
forkret_kthread(void)
{
  // 进程锁应已持有，释放之以允许并发
  struct proc *p = myproc();
  release(&p->lock);

  // 切入实际入口（不返回）
  kthread_entry();
  // 不可达
}

// 创建内核线程进程
int
create_process(void (*entry)(void))
{
  struct proc *p = alloc_process();
  if (p == 0)
    return -1;

  // 设置父子关系需持有 wait_lock，避免与 wait 的竞态
  acquire(&wait_lock);
  acquire(&p->lock);
  p->parent = myproc();                 // 设定父进程
  p->entry = entry;
  safestrcpy(p->name, "kproc", sizeof(p->name));
  p->state = RUNNABLE;
  release(&p->lock);
  release(&wait_lock);

  return p->pid;
}

// 主动让出 CPU
void
yield(void)
{
  struct proc *p = myproc();
  acquire(&p->lock);
  p->state = RUNNABLE;
  sched();
  release(&p->lock);
}

// 切回 scheduler 的核心原语
void
sched(void)
{
  struct proc *p = myproc();
  struct cpu *c = mycpu();

  if (!holding_spin(&p->lock))
    panic("sched p->lock");
  if (c->noff != 1)
    panic("sched locks");
  if (p->state == RUNNING)
    panic("sched running");

  // 记录中断使能状态
  int intena = c->intena;

  swtch(&p->context, &c->context);

  // 返回到该进程继续执行
  c->intena = intena;
}

// 终止当前进程
void
exit_process(int status)
{
  struct proc *p = myproc();

  // 关闭资源（文件、cwd 等，当前内核线程不涉及）

  // 处理子进程：孤儿交给“init”，此处简化为无“init”
  acquire(&wait_lock);
  for (struct proc *pp = proc; pp < &proc[NPROC]; pp++) {
    if (pp->parent == p) {
      // 交给无父进程（可扩展为 initproc）
      pp->parent = 0;
      // 唤醒可能等待的 initproc
    }
  }

  // 唤醒父进程
  acquire(&p->lock);
  p->xstate = status;
  p->state = ZOMBIE;
  release(&wait_lock);

  // 唤醒父进程：在 wait() 中睡眠于 &wait_lock
  wakeup(p->parent);

  sched();
  panic("zombie exit"); // 不可达
}

// 等待子进程退出
int
wait_process(int *status)
{
  struct proc *p = myproc();

  acquire(&wait_lock);
  for (;;) {
    int havekids = 0;
    for (struct proc *pp = proc; pp < &proc[NPROC]; pp++) {
      if (pp->parent == p) {
        // 找到子进程
        havekids = 1;
        acquire(&pp->lock);
        if (pp->state == ZOMBIE) {
          int pid = pp->pid;
          if (status) {
            *status = pp->xstate;
          }
          free_process(pp);
          release(&pp->lock);
          release(&wait_lock);
          return pid;
        }
        release(&pp->lock);
      }
    }

    if (!havekids) {
      release(&wait_lock);
      return -1;
    }

    // 睡眠等待子进程状态变化
    sleep(p, &wait_lock);
  }
}

// 睡眠/唤醒
void
sleep(void *chan, struct spinlock *lk)
{
  struct proc *p = myproc();

  // 持有 p->lock 以原子地设置睡眠状态
  acquire(&p->lock);
  release(lk);

  p->chan = chan;
  p->state = SLEEPING;

  sched();

  // 被唤醒后
  p->chan = 0;

  release(&p->lock);
  acquire(lk);
}

void
wakeup(void *chan)
{
  for (struct proc *p = proc; p < &proc[NPROC]; p++) {
    if (p == myproc()) continue;
    acquire(&p->lock);
    if (p->state == SLEEPING && p->chan == chan) {
      p->state = RUNNABLE;
    }
    release(&p->lock);
  }
}

// 调度器：轮转遍历进程表
void
scheduler(void)
{
  struct cpu *c = mycpu();
  c->proc = 0;

  for (;;) {
    intr_on(); // 开启中断，允许设备中断

    for (struct proc *p = proc; p < &proc[NPROC]; p++) {
      acquire(&p->lock);
      if (p->state == RUNNABLE) {
        p->state = RUNNING;
        c->proc = p;
        swtch(&c->context, &p->context);
        // 切回调度器
        c->proc = 0;
      }
      release(&p->lock);

      // 关键：在释放 p->lock（pop_off 完成）之后再开启中断，
      // 避免 pop_off 在中断已开时触发 "pop_off - interruptible"。
      intr_on();
    }
  }
}

// 提供强符号覆盖 trap.c 中的弱实现，使时钟中断可以触发抢占
// 当有正在运行的进程时，触发一次 yield 来实现抢占式调度。
void
scheduler_tick(void)
{
// no-op
}